from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from datetime import datetime, timedelta
import ipaddress


app = Flask(__name__)

# 允许跨域请求
CORS(app)

@app.route("/ca", methods=["POST"])
def get_ca_certificate():
    data = request.json
    dns_name = data.get('dnsNames', [])
    ip_addresses = data.get('ipAddresses', [])
    # 验证域名和IP地址是否合法
    # for dns_name in dns_name:
    #     if not dns_name.endswith('.test.net'):
    #         return jsonify({'error': 'Invalid DNS name'}), 400
    for ip_address in ip_addresses:
        try:
            ipaddress.ip_address(ip_address)
        except ValueError:
            return jsonify({'error': 'Invalid IP address'}), 400
        
    # 生成CA的私钥
    ca_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )

    # 生成CA证书
    ca_name = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"My CA Company"),
        x509.NameAttribute(NameOID.COMMON_NAME, u"My CA Root"),
    ])

    ca_certificate = (
        x509.CertificateBuilder()
        .subject_name(ca_name)
        .issuer_name(ca_name)
        .public_key(ca_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.utcnow())
        .not_valid_after(datetime.utcnow() + timedelta(days=36500))  # 证书有效期10年
        .add_extension(
            x509.BasicConstraints(ca=True, path_length=None), critical=True,
        )
        .sign(ca_key, hashes.SHA256())
    )

    # 保存CA的私钥和证书
    # with open("ca_private_key.pem", "wb") as f:
    #     f.write(ca_key.private_bytes(
    #         encoding=serialization.Encoding.PEM,
    #         format=serialization.PrivateFormat.TraditionalOpenSSL,
    #         encryption_algorithm=serialization.NoEncryption(),
    #     ))

    ca_key_pem = ca_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )

    # with open("ca_certificate.pem", "wb") as f:
    #     f.write(ca_certificate.public_bytes(serialization.Encoding.PEM))

    ca_certificate_pem = ca_certificate.public_bytes(serialization.Encoding.PEM)
    
    # 生成服务证书的私钥
    service_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )

    # 定义服务证书的详细信息
    service_name = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"My Service Company"),
        x509.NameAttribute(NameOID.COMMON_NAME, u"mydomain.com"),
    ])

    # 创建SubjectAlternativeName扩展
    alt_names = x509.SubjectAlternativeName([
        x509.DNSName(dns_name) for dns_name in dns_name
    ] + [
        x509.IPAddress(ipaddress.ip_address(ip_address)) for ip_address in ip_addresses
    
        # x509.DNSName(u"*.test.net"),
        # x509.IPAddress(ipaddress.ip_address(u"192.168.102.210")),
        # x509.DNSName(u"anotherdomain.com"),
        # x509.IPAddress(ipaddress.ip_address(u"192.168.31.1"))
    ])

    service_certificate = (
        x509.CertificateBuilder()
        .subject_name(service_name)
        .issuer_name(ca_certificate.subject)
        .public_key(service_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.utcnow())
        .not_valid_after(datetime.utcnow() + timedelta(days=3650))  # 证书有效期10年
        .add_extension(alt_names, critical=False)
        .sign(ca_key, hashes.SHA256())
    )

    # 保存服务证书和私钥
    # with open("service_private_key.pem", "wb") as f:
    #     f.write(service_key.private_bytes(
    #         encoding=serialization.Encoding.PEM,
    #         format=serialization.PrivateFormat.TraditionalOpenSSL,
    #         encryption_algorithm=serialization.NoEncryption(),
    #     ))

    service_key_pem = service_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )

    # with open("service_certificate.pem", "wb") as f:
    #     f.write(service_certificate.public_bytes(serialization.Encoding.PEM))

    service_certificate_pem = service_certificate.public_bytes(serialization.Encoding.PEM)
    
    return jsonify({
        'ca_key': ca_key_pem.decode('utf-8'),
        'ca_certificate': ca_certificate_pem.decode('utf-8'),
        'service_key': service_key_pem.decode('utf-8'),
        'service_certificate': service_certificate_pem.decode('utf-8'),
    }), 200

if __name__ == "__main__":
    app.run(debug=True, host='0.0.0.0', port=5000)